"""
Author: Morphlng
Date: 2024-03-19 13:50:34
LastEditTime: 2024-03-19 13:50:34
LastEditors: Morphlng
Description: Test for vector3d
FilePath: /DrivingGym/tests/common/geom/vector3d_test.py
"""

import math

import numpy as np

from driving_gym.simulation.common.geom import Vector2D, Vector3D


def is_close(a, b):
    return math.isclose(a + 1, b + 1)


class TestVector3D:

    def test_initialization(self):
        vec = Vector3D(1, 2, 3)
        assert (
            vec.x == 1 and vec.y == 2 and vec.z == 3
        ), "Vector3D initialization failed."

    def test_as_numpy_array(self):
        vec = Vector3D(1, 2, 3)
        assert np.array_equal(
            vec.as_numpy_array(), np.array([1, 2, 3])
        ), "Conversion to numpy array failed."

    def test_as_numpy_array_2D(self):
        vec = Vector3D(1, 2, 3)
        assert np.array_equal(
            vec.as_numpy_array_2D(), np.array([1, 2])
        ), "Conversion to 2D numpy array failed."

    def test_as_vector_2D(self):
        vec = Vector3D(1, 2, 3)
        vec2d = vec.as_vector_2D()
        assert (
            isinstance(vec2d, Vector2D) and vec2d.x == 1 and vec2d.y == 2
        ), "Conversion to Vector2D failed."

    def test_cross(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(4, 5, 6)
        cross_product = vec1.cross(vec2)
        assert isinstance(cross_product, Vector3D), "Cross product type incorrect."
        assert (
            cross_product.x == -3 and cross_product.y == 6 and cross_product.z == -3
        ), "Cross product computation failed."

    def test_distance(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(4, 5, 6)
        assert is_close(
            vec1.distance(vec2), math.sqrt(27)
        ), "Euclidean distance computation failed."

    def test_distance_2d(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(4, 5, 6)
        assert is_close(
            vec1.distance_2d(vec2), math.sqrt(18)
        ), "2D distance computation failed."

    def test_l1_distance(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(4, 5, 6)
        assert vec1.l1_distance(vec2) == 9, "L1 distance computation failed."

    def test_l2_distance(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(4, 5, 6)
        assert is_close(
            vec1.l2_distance(vec2), math.sqrt(27)
        ), "L2 distance computation failed."

    def test_distance_to_polygon(self):
        vec = Vector3D(0, 2, 0)
        polygon = [(1, 1), (1, -1), (-1, -1), (-1, 1)]
        assert is_close(
            vec.distance_to_polygon(polygon), 1.0
        ), "Distance to polygon computation failed."

        vec = Vector3D(0, 0, 0)
        assert is_close(
            vec.distance_to_polygon(polygon), 0.0
        ), "Point inside polygon distance computation failed."

        vec = Vector3D(2, 2, 0)
        assert is_close(
            vec.distance_to_polygon(polygon), math.sqrt(2)
        ), "Distance to polygon computation failed."

    def test_magnitude(self):
        vec = Vector3D(3, 4, 0)
        assert is_close(vec.magnitude(), 5), "Magnitude computation failed."

    def test_length(self):
        vec = Vector3D(3, 4, 0)
        assert is_close(vec.length(), 5), "Length computation failed."

    def test_make_unit_vector(self):
        vec = Vector3D(3, 4, 0)
        unit_vec = vec.make_unit_vector()
        assert is_close(unit_vec.magnitude(), 1), "Unit vector computation failed."

    def test_rotate(self):
        vec = Vector3D(1, 0, 0)

        try:
            import random

            from scipy.spatial.transform import Rotation as R
        except ImportError:
            # Simple test only
            rotated_vec = vec.rotate(90)
            assert (
                is_close(rotated_vec.x, 0)
                and is_close(rotated_vec.y, 1)
                and is_close(rotated_vec.z, 0)
            ), "Rotation failed."
            return

        angle = random.uniform(0, 360)
        axis = np.random.rand(3)

        # Use scipy to rotate
        r = R.from_rotvec(np.deg2rad(angle) * axis / np.linalg.norm(axis))
        scipy_rotated_vector = r.apply(vec.as_numpy_array())

        # Use our method to rotate
        rotated_vec = vec.rotate(angle, axis)

        assert np.allclose(
            scipy_rotated_vector, rotated_vec.as_numpy_array()
        ), "Rotation failed."

    def test_copy(self):
        vec = Vector3D(1, 2, 3)
        copied_vec = vec.copy()
        assert copied_vec == vec and copied_vec is not vec, "Copy failed."

    def test_abs(self):
        vec = Vector3D(-1, -2, -3)
        abs_vec = abs(vec)
        assert (
            abs_vec.x == 1 and abs_vec.y == 2 and abs_vec.z == 3
        ), "Absolute value failed."

    def test_add(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(4, 5, 6)
        sum_vec = vec1 + vec2
        assert sum_vec.x == 5 and sum_vec.y == 7 and sum_vec.z == 9, "Addition failed."

    def test_sub(self):
        vec1 = Vector3D(4, 6, 8)
        vec2 = Vector3D(1, 2, 3)
        diff_vec = vec1 - vec2
        assert (
            diff_vec.x == 3 and diff_vec.y == 4 and diff_vec.z == 5
        ), "Subtraction failed."

    def test_mul(self):
        vec = Vector3D(1, 2, 3)
        scaled_vec = vec * 3
        assert (
            scaled_vec.x == 3 and scaled_vec.y == 6 and scaled_vec.z == 9
        ), "Scalar multiplication failed."

    def test_truediv(self):
        vec = Vector3D(3, 6, 9)
        divided_vec = vec / 3
        assert (
            divided_vec.x == 1 and divided_vec.y == 2 and divided_vec.z == 3
        ), "Scalar division failed."

    def test_eq(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(1, 2, 3)
        assert vec1 == vec2, "Equality check failed."

    def test_ne(self):
        vec1 = Vector3D(1, 2, 3)
        vec2 = Vector3D(3, 2, 1)
        assert vec1 != vec2, "Inequality check failed."
